from typing import Optional

import aiohttp
import time

from bot.insurance_planner_gpt.qw_server.common import wework_constants
from bot.insurance_planner_gpt.qw_server.utils import common_util


class AccessTokenManager:
    def __init__(self):
        self.company_tokens = {}

    async def get_access_token_by_company_id(self, company_id: str) -> Optional[str]:
        token_info = self.company_tokens.get(company_id)
        if token_info and time.time() < token_info['expires_at']:
            return token_info['access_token']
        else:
            return await self._refresh_access_token(company_id)

    async def get_access_token_by_app_key(self, app_key: str) -> Optional[str]:
        for company_id, token_info in self.company_tokens.items():
            if token_info['app_key'] == app_key and time.time() < token_info['expires_at']:
                return token_info['access_token']
        company_id = await self._get_company_id_by_app_key(app_key)
        if company_id:
            return await self._refresh_access_token(company_id)
        else:
            return None

    async def _get_company_id_by_app_key(self, app_key: str) -> Optional[str]:
        for company_info in common_util.wework_config_list:
            if company_info['WE_WORK_APP_KEY'] == app_key:
                return company_info['COMPANY_ID']
        return None

    async def _refresh_access_token(self, company_id: str) -> Optional[str]:
        company_info = common_util.get_info_by_company_id(common_util.wework_config_list, company_id)
        app_key = company_info.get('WE_WORK_APP_KEY')
        app_secret = company_info.get('WE_WORK_APP_SECRET')
        payload = {
            "app_key": app_key,
            "app_secret": app_secret,
        }
        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(wework_constants.WEWORK_INTERFACE_GET_ACCESS_TOKEN, json=payload) as response:
                    data = await response.json()
                    if response.status == 200 and data.get("errcode") == 0:
                        access_token = data["data"]["access_token"]
                        expires_at = time.time() + data["data"]["expires_in"]
                        self.company_tokens[company_id] = {
                            'app_key': app_key,
                            'access_token': access_token,
                            'expires_at': expires_at
                        }
                        return access_token
                    else:
                        print(f"获取token异常: {data.get('errmsg', '未知错误')}")
        except Exception as e:
            print(f"Exception: {str(e)}")
            return None


token_manager = AccessTokenManager()

# 替换为实际的 app_key 和 app_secret
# app_key = settings.WE_WORK_APP_KEY
# app_secret = settings.WE_WORK_APP_SECRET
#
# token_manager = AccessTokenManager(app_key, app_secret)
#
# # 使用 Access Token
# access_token = token_manager.get_access_token()
# print(f"Current Access Token: {access_token}")
